import os
import pandas as pd
import pymysql
from sqlalchemy import create_engine, text, inspect
from tqdm import tqdm
import logging


# ---------------------- 配置 ----------------------
class DBConfig:
    HOST = "127.0.0.1"  # 根据实际修改
    USER = "root"
    PASSWD = "root"
    DB = "tushare"
    PORT = 3306
    CHARSET = "utf8mb4"


SQL_QUERY = text("""
SELECT 
    d.trade_date, d.close AS price, 
    m.buy_elg_vol, m.sell_elg_vol,
    i.i_close AS index_price, i.i_vol AS index_vol
FROM daily_data d
JOIN moneyflow_data m ON d.ts_code=m.ts_code AND d.trade_date=m.trade_date
JOIN index_data i ON d.trade_date=i.trade_date
WHERE d.ts_code=:ts_code 
  AND d.trade_date BETWEEN :start_date AND :end_date
""")

PARAMS = {
    "ts_code": "000001.SH",  # 示例股票代码
    "start_date": "2020-01-01",
    "end_date": "2023-12-31"
}


class DataFetcher:
    def __init__(self):
        self.engine = create_engine(
            f"mysql+pymysql://{DBConfig.USER}:{DBConfig.PASSWD}@"
            f"{DBConfig.HOST}:{DBConfig.PORT}/{DBConfig.DB}",
            pool_size=5
        )
        self.logger = self._init_logger()

    def _init_logger(self):
        logger = logging.getLogger(__name__)
        logger.setLevel(logging.INFO)
        formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')

        # 文件日志
        fh = logging.FileHandler('data_fetch.log')
        fh.setFormatter(formatter)
        logger.addHandler(fh)

        # 控制台日志
        ch = logging.StreamHandler()
        ch.setFormatter(formatter)
        logger.addHandler(ch)
        return logger

    def _validate_schema(self):
        """校验表结构和字段完整性"""
        required_tables = {'daily_data', 'moneyflow_data', 'index_data'}
        inspector = inspect(self.engine)

        # 检查表是否存在
        existing_tables = set(inspector.get_table_names())
        missing_tables = required_tables - existing_tables
        if missing_tables:
            self.logger.error(f"缺失表: {missing_tables}")
            raise ValueError("数据库表结构不完整")

        # 检查字段是否存在
        required_cols = {
            'daily_data': ['trade_date', 'ts_code', 'close'],
            'moneyflow_data': ['buy_elg_vol', 'sell_elg_vol']
        }
        for table, cols in required_cols.items():
            existing_cols = {c['name'] for c in inspector.get_columns(table)}
            missing = set(cols) - existing_cols
            if missing:
                raise KeyError(f"表 {table} 缺失字段: {missing}")

    def fetch(self):
        """主获取流程"""
        try:
            self._validate_schema()

            with self.engine.connect() as conn:
                # 分页查询（每次获取30天）
                date_ranges = pd.date_range(
                    start=PARAMS['start_date'],
                    end=PARAMS['end_date'],
                    freq='30D'
                ).tolist()

                dfs = []
                for i in tqdm(range(len(date_ranges) - 1)):
                    start = date_ranges[i].strftime('%Y-%m-%d')
                    end = date_ranges[i + 1].strftime('%Y-%m-%d')

                    df_chunk = pd.read_sql(
                        SQL_QUERY.bindparams(
                            ts_code=PARAMS['ts_code'],
                            start_date=start,
                            end_date=end
                        ),
                        conn,
                        parse_dates=['trade_date']
                    )
                    dfs.append(df_chunk)
                    self.logger.info(f"已获取 {start} 至 {end} 数据，共 {len(df_chunk)} 条")

                full_df = pd.concat(dfs)
                full_df.to_csv("raw_data.csv", index=False)
                self.logger.info(f"数据获取完成，总计 {len(full_df)} 条")

        except Exception as e:
            self.logger.error(f"获取失败: {str(e)}", exc_info=True)
            raise


if __name__ == "__main__":
    fetcher = DataFetcher()
    fetcher.fetch()
